package com.net.shard;


import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**

 */
public class ShardDataSource extends AbstractRoutingDataSource {
    private static final Logger LOGGER = LoggerFactory.getLogger(ShardDataSource.class);

    private  ShardDataSourceConfigurer shardDataSourceConfigurer;

    //增加一个定时器 定时扫描新增的数据源
    private final ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1,
            new ThreadFactoryBuilder().setDaemon(true).setNameFormat("data-source-schedule").build());


    public ShardDataSourceConfigurer getShardDataSourceConfigurer() {
        return shardDataSourceConfigurer;
    }

    public void setShardDataSourceConfigurer(ShardDataSourceConfigurer shardDataSourceConfigurer) {
        this.shardDataSourceConfigurer = shardDataSourceConfigurer;
    }


    @Override
    protected Object determineCurrentLookupKey() {
        ShardContext context = ShardContextHolder.getShardContext();
        if(context == null){
            throw new RuntimeException("Shard Context is null for |" + shardDataSourceConfigurer.getDataSourceName());
        }
        return context.getShardId();
    }

    @Override
    public void afterPropertiesSet() {
        Map<Long, DataSource> dataSourceMap = shardDataSourceConfigurer.initDataSources();
        Map<Object, Object> objectMap = new HashMap<>();
        objectMap.putAll(dataSourceMap);

        super.setTargetDataSources(objectMap);
        super.setLenientFallback(false);  //不需要用默认的DS
        super.afterPropertiesSet(); //需要初始化一下

        //下次在1分钟之后调度
        executor.scheduleWithFixedDelay(this::refreshDataSource, 5, 5, TimeUnit.MINUTES);
    }

    //刷新数据源
    private void refreshDataSource() {
        try {
            //检查下数据库数量变了没有
            if (!shardDataSourceConfigurer.needRefreshDataSource()) {
                return;
            }

            Map<Long, DataSource> dataSourceMap = shardDataSourceConfigurer.refreshDataSources();
            Map<Object, Object> objectMap = new HashMap<>();
            objectMap.putAll(dataSourceMap);

            super.setTargetDataSources(objectMap);
            super.afterPropertiesSet(); //刷新配置
        } catch (Exception e) {
            LOGGER.error("refreshDataSource error", e);
        }
    }
}
